# scheduler/sync_templates.py
import json
import logging

from reportagentic.core.db import TemplateRepository, MetadataRepository
from reportagentic.core.llm import LLMClient
from reportagentic.core.vector import VectorClient


class TemplateScheduler:

    def __init__(self, template_repo: TemplateRepository,
                 meta_repo: MetadataRepository,
                 llm_client: LLMClient,
                 vector_client: VectorClient):
        self.template_repo = template_repo
        self.meta_repo = meta_repo
        self.llm_client = llm_client
        self.vector_client = vector_client

    def sync_all_templates_to_vector_store(self):
        logging.info("Starting template vector synchronization task...")

        try:
            # 1. 从数据库获取所有激活的模板
            active_templates = self.template_repo.get_all_active_templates()
            if not active_templates:
                logging.info("No active templates found to sync. Task finished.")
                return

            logging.info(f"Found {len(active_templates)} active templates to process.")

            # 2. 准备批量处理的数据
            ids_to_upsert = []
            texts_to_embed = []
            payloads_to_upsert = []

            for template in active_templates:
                # 组合所有文本信息用于生成一个综合的向量
                name = template.get('template_name', '')
                description = template.get('description', '')

                # example_descriptions 在数据库中是JSON字符串，需要解析
                try:
                    example_descs_str = template.get('example_descriptions')
                    example_descs_list = json.loads(example_descs_str) if example_descs_str else []
                except (json.JSONDecodeError, TypeError):
                    example_descs_list = []

                # 构建用于 embedding 的文本块
                text_parts = [f"报表模板: {name}", f"功能描述: {description}"]
                if example_descs_list:
                    text_parts.append("用例: " + " | ".join(example_descs_list))

                text_to_embed = "\n".join(text_parts)

                # 添加到批量处理列表中
                ids_to_upsert.append(template['id'])
                texts_to_embed.append(text_to_embed)
                payloads_to_upsert.append({"name": name, "description": description})

            # 3. 批量生成向量
            logging.info(f"Generating embeddings for {len(texts_to_embed)} templates in a single batch...")
            vectors = self.llm_client.get_embeddings(texts_to_embed)

            if not vectors or len(vectors) != len(ids_to_upsert):
                logging.error("Failed to generate embeddings or mismatched count. Aborting sync.")
                return

            # 4. 批量写入（或替换）到向量存储
            self.vector_client.upsert_vectors(
                collection_name="template_vectors",
                ids=ids_to_upsert,
                vectors=vectors,
                payloads=payloads_to_upsert
            )
            logging.info(f"Successfully upserted {len(ids_to_upsert)} vectors into 'template_vectors' collection.")
        except Exception as e:
            logging.error(f"An error occurred during the synchronization task: {e}", exc_info=True)
